# -*- coding: utf-8 -*-
"""
文件名: src/utils/loc_enhance.py

功能概述:
    针对包含坐标标注(检测框/定位点)的数据，进行图像增强并同步变换归一化坐标:
      - 支持操作: 'h' (水平翻转), 'v' (垂直翻转), 'vh' (水平+垂直), 'rot90' (顺时针90°)
      - 输入和输出的坐标均为归一化坐标:
            检测:   [x1, y1, x2, y2], 0~1 浮点数，保证 x1<x2, y1<y2
            定位点: [x, y], 0~1 浮点数
      - 保存增强后图片到 save_dir/images 下, 文件名在原名后追加后缀(_h/_v/_vh/_r90).png
      - 生成与原始结构一致的增强标注条目(仅包含坐标相关字段)，便于主流程合并

使用样例:
    from src.utils.loc_enhance import CoordImageEnhancer

    enhancer = CoordImageEnhancer(
        save_dir="./experiments/dataset_enhanced/coor_qa",
        count_strategy="random",
        per_image_k=2,
        op_probs={"h": 0.4, "v": 0.4, "vh": 0.15, "rot90": 0.05},
        seed=42
    )

    # 单图增强
    new_items = enhancer.enhance_one(
        image_path="/abs/path/to/img.png",
        image_name="dataset_xxx_0001.png",
        # 检测: 中文标签 -> 多个归一化框
        detection_dict={"视盘": [[0.45,0.42,0.52,0.50]]},
        # 定位: 中文标签 -> 多个归一化点
        location_dict={"黄斑": [[0.53,0.48]]},
        source_dataset="Messidor",
        label=None,            # 若有分类文字可传入, 否则传 None
        quality=None           # 若有质量文字可传入, 否则传 None
    )

    # 批量增强
    batch_items = enhancer.enhance_batch(list_of_records)
    # 其中 list_of_records 的每个元素可形如:
    # {
    #   "image_path": "...",
    #   "image_name": "...",
    #   "detection": {...},         # 可选
    #   "location": {...},      # 可选
    #   "source_dataset": "xxx",    # 可选
    #   "label": ["中文病名1", ...]  # 可选(分类文本)
    #   "quality": "中文质量句子"     # 可选(质量文本)
    # }

注意:
    - 本模块聚焦“坐标类增强”，不涉及分割 mask 的裁剪或重采样。
    - 所有异常一律 raise，并在异常信息中附带关键上下文以便调试。
    - 若同一图片被多次增强, 每个操作会生成一个独立样本条目并各自保存图像。
    - 由于计算坐标相对复杂，我们只采用三种增强方式

作者: zym1105
时间: 2025-10-10
"""

import os
import io
import json
import math
import copy
import random
from typing import Dict, List, Tuple, Optional, Any

from PIL import Image
from PIL.Image import Image as PILImage

# -----------------------------
# 基础几何变换 (作用于归一化坐标)
# -----------------------------

def _flip_h_bbox_norm(b: List[float]) -> List[float]:
    x1, y1, x2, y2 = b
    nx1, nx2 = 1.0 - x2, 1.0 - x1
    ny1, ny2 = y1, y2
    nx1, nx2 = min(nx1, nx2), max(nx1, nx2)
    ny1, ny2 = min(ny1, ny2), max(ny1, ny2)
    return [nx1, ny1, nx2, ny2]

def _flip_v_bbox_norm(b: List[float]) -> List[float]:
    x1, y1, x2, y2 = b
    nx1, nx2 = x1, x2
    ny1, ny2 = 1.0 - y2, 1.0 - y1
    nx1, nx2 = min(nx1, nx2), max(nx1, nx2)
    ny1, ny2 = min(ny1, ny2), max(ny1, ny2)
    return [nx1, ny1, nx2, ny2]

def _rot90_bbox_norm(b, clockwise: bool = True):
    """归一化坐标的 bbox 旋转90°。
    clockwise=True  -> 视觉上顺时针（PIL.ROTATE_270）
    clockwise=False -> 视觉上逆时针（PIL.ROTATE_90）
    """
    x1, y1, x2, y2 = b
    pts = [(x1,y1), (x1,y2), (x2,y1), (x2,y2)]

    if clockwise:
        # 视觉顺时针：与 PIL.ROTATE_270 一致
        rot = [(1 - y, x) for (x, y) in pts]
    else:
        # 视觉逆时针：与 PIL.ROTATE_90 一致
        rot = [(y, 1 - x) for (x, y) in pts]

    xs = [p[0] for p in rot]
    ys = [p[1] for p in rot]
    return [min(xs), min(ys), max(xs), max(ys)]


def _flip_h_point_norm(p: List[float]) -> List[float]:
    x, y = p
    return [1.0 - x, y]

def _flip_v_point_norm(p: List[float]) -> List[float]:
    x, y = p
    return [x, 1.0 - y]

def _rot90_point_norm(p):
    # 视觉上顺时针 90° 对齐 PIL.ROTATE_270
    x, y = p
    return [1.0 - y, x]

def _apply_bbox_op(b: List[float], op: str) -> List[float]:
    if op == "h":
        return _flip_h_bbox_norm(b)
    if op == "v":
        return _flip_v_bbox_norm(b)
    if op == "vh":  # 水平+垂直 == 180° 旋转
        return _flip_v_bbox_norm(_flip_h_bbox_norm(b))
    if op == "rot90":
        return _rot90_bbox_norm(b)
    raise ValueError(f"未知操作: {op}")

def _apply_point_op(p: List[float], op: str) -> List[float]:
    if op == "h":
        return _flip_h_point_norm(p)
    if op == "v":
        return _flip_v_point_norm(p)
    if op == "vh":
        return _flip_v_point_norm(_flip_h_point_norm(p))
    if op == "rot90":
        return _rot90_point_norm(p)
    raise ValueError(f"未知操作: {op}")


# -----------------------------
# 图像变换
# -----------------------------

def _apply_image_op(img: PILImage, op: str) -> PILImage:
    if op == "h":
        return img.transpose(Image.FLIP_LEFT_RIGHT)
    if op == "v":
        return img.transpose(Image.FLIP_TOP_BOTTOM)
    if op == "vh":
        # 先水平后垂直，与 180° 旋转等价
        return img.transpose(Image.FLIP_LEFT_RIGHT).transpose(Image.FLIP_TOP_BOTTOM)
    if op == "rot90":
        # 顺时针 90°
        return img.transpose(Image.ROTATE_270)
    raise ValueError(f"未知操作: {op}")
# -----------------------------
# 采样操作集合 (支持多操作/不等概率)
# -----------------------------

def _weighted_sample_ops(
    ops: List[str],
    weights: List[float],
    k: int,
    rng: random.Random
) -> List[str]:
    """
    带权不放回采样 k 个操作。若 k > len(ops) 则上限为 len(ops)。
    """
    if k <= 0:
        return []
    k = min(k, len(ops))
    # 简单重复抽取 + 去重，直到取满或无可选
    selected = []
    pool = list(zip(ops, weights))
    while len(selected) < k and pool:
        # 归一化
        tot = sum(max(w, 0.0) for _, w in pool)
        if tot <= 0:
            # 无有效权重则均匀
            idx = rng.randrange(len(pool))
        else:
            r = rng.random() * tot
            acc = 0.0
            idx = 0
            for i, (_, w) in enumerate(pool):
                acc += max(w, 0.0)
                if r <= acc:
                    idx = i
                    break
        op, _ = pool.pop(idx)
        selected.append(op)
    return selected


# -----------------------------
# 主类: 坐标增强器
# -----------------------------

class CoordImageEnhancer:
    """
    对坐标类样本进行增强:
      - 保存增强后的图片
      - 同步更新 detection/location 的归一化坐标
      - 返回增强条目列表
    """

    def __init__(
        self,
        save_dir: str,
        count_strategy: str = "random",      # "random" 或 "fixed"
        per_image_k: int = 2,               # fixed 时每图增强张数
        op_probs: Optional[Dict[str, float]] = None,
        seed: int = 42
    ):
        """
        参数:
            save_dir:          增强数据集根目录(会写入 images/ )
            count_strategy:    'random' -> 每图随机增强(1~3)；'fixed' -> 固定 per_image_k
            per_image_k:       fixed 模式下每图增强张数
            op_probs:          操作概率, 默认 {'h':0.4,'v':0.4,'vh':0.15,'rot90':0.05}
            seed:              随机种子
        """
        self.save_dir = os.path.abspath(save_dir)
        self.images_dir = os.path.join(self.save_dir, "images")
        os.makedirs(self.images_dir, exist_ok=True)

        self.count_strategy = count_strategy
        self.per_image_k = max(1, int(per_image_k))
        self.rng = random.Random(seed)

        self.ops = ["h", "v", "vh", "rot90"]
        default_probs = {"h": 0.4, "v": 0.4, "vh": 0.15, "rot90": 0.05}
        self.op_probs = op_probs or default_probs
        # 保证存在所有 key
        for k in self.ops:
            self.op_probs.setdefault(k, default_probs[k])

    # ---------- 外部 API ----------

    def enhance_one(
        self,
        image_path: str,
        image_name: str,
        detection_dict: Optional[Dict[str, List[List[float]]]] = None,
        location_dict: Optional[Dict[str, List[List[float]]]] = None,
        source_dataset: Optional[str] = None,
        label: Optional[List[str]] = None,     # 分类中文标签列表(可选)
        quality: Optional[str] = None          # 质量中文描述(可选)
    ) -> List[Dict[str, Any]]:
        """
        对单张图片执行若干增强操作, 返回增强后的样本条目列表(每个操作对应一条)。

        返回的条目示例:
        {
            "image_name": "orig_xxx_h.png",
            "image_path":  "/.../images/orig_xxx_h.png",
            "source_image_name": "orig_xxx.png",
            "source_dataset": "Messidor",
            "ops": "h",  # 本次增强操作
            "diagnosis": {
                "classification": {"text_cn": ["..."], "quality_cn": "..."}   # 可选
            },
            "detection": {
                "视盘": [[x1,y1,x2,y2], ...]    # 归一化
            },
            "location": {
                "黄斑": [[x,y], ...]            # 归一化
            }
        }
        """
        try:
            if not os.path.exists(image_path):
                raise FileNotFoundError(f"原图不存在: {image_path}")

            img = Image.open(image_path).convert("RGB")

            # 采样操作集合
            k = self._decide_k()
            ops = self._sample_ops(k)

            # 逐操作生成新样本
            new_items: List[Dict[str, Any]] = []
            base_stem, _ = os.path.splitext(image_name)

            for op in ops:
                out_img = _apply_image_op(img, op)

                # 保存文件名后缀
                suffix = self._suffix_of(op)
                new_name = f"{base_stem}{suffix}.png"
                out_path = os.path.join(self.images_dir, new_name)
                out_img.save(out_path)

                # 同步变换坐标 (归一化)
                new_det = self._transform_detection(detection_dict, op)
                new_loc = self._transform_location(location_dict, op)

                item = {
                    "image_name": new_name,
                    "image_path": out_path,
                    "source_image_name": image_name,
                    "ops": op
                }
                if source_dataset:
                    item["source_dataset"] = source_dataset
                # 可选: 分类/质量信息在增强样本中保留下来(训练时可用)

                if label:
                    item['text_cn']=list[label]
                if quality:
                    item['quality']=quality
                if new_det:
                    item["detection"] = new_det
                if new_loc:
                    item["location"] = new_loc

                new_items.append(item)

            return new_items

        except Exception as e:
            raise RuntimeError(
                f"[enhance_one] 失败: image_name={image_name}, image_path={image_path}. "
                f"原因: {repr(e)}"
            ) from e

    # ---------- 工具/内部 ----------

    def _decide_k(self) -> int:
        if self.count_strategy == "fixed":
            return self.per_image_k
        # random: 均匀采样 1~3
        return self.rng.choice([1, 2, 3])

    def _sample_ops(self, k: int) -> List[str]:
        """
        不放回带权采样 k 个操作。
        注意: 我们不区分 h 与 v 的重要性，默认概率相当；'vh' 代表联合翻转；
             这样可以减少单一方向的样式偏置。
        """
        ops = ["h", "v", "vh", "rot90"]
        weights = [self.op_probs.get("h", 0.4),
                   self.op_probs.get("v", 0.4),
                   self.op_probs.get("vh", 0.15),
                   self.op_probs.get("rot90", 0.05)]
        return _weighted_sample_ops(ops, weights, k, self.rng)

    @staticmethod
    def _suffix_of(op: str) -> str:
        return {
            "h": "_h",
            "v": "_v",
            "vh": "_vh",
            "rot90": "_r90"
        }[op]

    @staticmethod
    def _transform_detection(
        detection_dict: Optional[Dict[str, List[List[float]]]],
        op: str
    ) -> Optional[Dict[str, List[List[float]]]]:
        if not detection_dict:
            return None
        out: Dict[str, List[List[float]]] = {}
        for k, boxes in detection_dict.items():
            new_boxes = []
            for b in boxes:
                if not isinstance(b, (list, tuple)) or len(b) != 4:
                    raise ValueError(f"检测框格式不合法: {b}")
                bb = [float(b[0]), float(b[1]), float(b[2]), float(b[3])]
                new_boxes.append(_apply_bbox_op(bb, op))
            out[k] = new_boxes
        return out

    @staticmethod
    def _transform_location(
        location_dict: Optional[Dict[str, List[List[float]]]],
        op: str
    ) -> Optional[Dict[str, List[List[float]]]]:
        if not location_dict:
            return None
        out: Dict[str, List[List[float]]] = {}
        for k, points in location_dict.items():
            new_pts = []
            for p in points:
                if not isinstance(p, (list, tuple)) or len(p) != 2:
                    raise ValueError(f"定位点格式不合法: {p}")
                pp = [float(p[0]), float(p[1])]
                new_pts.append(_apply_point_op(pp, op))
            out[k] = new_pts
        return out
